package nmf;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.PriorityQueue;

import org.ejml.ops.NormOps;
import org.ejml.simple.SimpleMatrix;

import utils.Summarizer;

/***
 * 
 * Uses the most important words in V to fill out H, and then back solves for
 * W.
 *
 */
public class TopWords extends NMFTool {
   
   private final int trainClass[], numberOfCategories;
   private Integer[] topWords;
   
   public TopWords(final int[] trainClass) {
      this.trainClass = trainClass;
      numberOfCategories = getNumCategories(trainClass, trainClass.length);
   }
   
   /**
    * This method initializes the rows of H with the rows in V that are most
    * pertinent to the distinct categories. Each category is alloted an 
    * approximately equal number of words to fill the rows of H. The words 
    * chosen for a particular category are picked in descending order of their 
    * significance. The number of words per category is chosen based on the 
    * compressed dimension. W is then chosen as an approximation through the 
    * least squares method.
    */
   @Override
   public void initialize(final SimpleMatrix V, final SimpleMatrix W, final SimpleMatrix H, final int[] trainClass) {
      if (trainClass == null) {
         System.out.println("Actual classifications unknown, " +
               "using simple initialization.");
         final NMFTool backup = new SimpleNMFTool();
         backup.initialize(V, W, H, trainClass);
      }
      
      //Determine how many words per category we can have
      //to just barely overflow H with rows.
      final int wordsPerCategory = W.numCols() / numberOfCategories + 1;
      topWords = getTopWords(V, wordsPerCategory);
      
      //Copy as many of the relevant rows into H as possible. Truncation
      //of the least important words is possible.
      for (int ndx = 0; ndx < H.numRows(); ndx++) {
         copyRow(H, V, topWords[ndx], ndx);
      }
      
      //RandomMatrices.setRandom(W.getMatrix(), 0, 1, new Random());
      //W.set(.015);
      final SimpleMatrix forcedSquare = H.mult(H.transpose());
      if (forcedSquare.determinant() != 0) {
         W.set(V.mult(H.transpose()).mult(forcedSquare.invert()));
      }
      else {
         System.out.println("Initialization for W not possible.");
         W.set(epsilon);
      }
   }
   
   private void copyRow(final SimpleMatrix H, final SimpleMatrix V, final int rowNumV, final int rowNumH) {
      for (int col = 0; col < V.numCols(); col++) {
         H.set(rowNumH, col, V.get(rowNumV, col));
      }
   }
   
   /**
    * Searches through the data matrix V for a list of row numbers.
    * Each document from a given category is perused, and the words
    * with the highest tf-idf score are recorded. This is done for each
    * category, and each category has the highest 'wordsPerCategory'
    * words recorded. The integer array returned features the most
    * prominent words of each category first, then the second most
    * prominent from each category, and so on and so forth.
    * 
    * 
    * @param V
    * @param wordsPerCategory
    * @return An equal number of top rated words for each category as an array
    * of indices pointing to the corresponding row.
    */
   public Integer[] getTopWords(final SimpleMatrix V, final int wordsPerCategory) {
      final ArrayList<Integer[]> categoryTopWords = new ArrayList<Integer[]>();
      
      for (int i = 1; i <= numberOfCategories; i++) {
         categoryTopWords.add(topNRows(V, i, wordsPerCategory));
      }
      
      final Integer[] topWords = new Integer[numberOfCategories * wordsPerCategory];
      int depth = 0;
      int currentWord = 0;
      
      while (currentWord < numberOfCategories * wordsPerCategory) {
         for (int i = 0; i < numberOfCategories; i++) {
            topWords[currentWord++] = categoryTopWords.get(i)[depth];
         }
         depth++;
      }
      
      /*for (int i = 0; i < categoryTopWords.size(); i++) {
         System.arraycopy(categoryTopWords.get(i), 0, topWords, pos, categoryTopWords.get(i).length);
         pos += categoryTopWords.get(i).length;
      }
      
      Arrays.sort(topWords);*/
      
      return topWords;
   }
   
   public Integer[] getTopWords() {
      return topWords.clone();
   }
   
   //Returns the top n most important rows to a given category in V.
   public Integer[] topNRows(final SimpleMatrix V, final int category, final int n) {
      new HashMap<Integer, Double>();
      final PriorityQueue<Integer> topWordsAcrossDocs = new PriorityQueue<Integer>();
      
      for (int i = 0; i < numberOfCategories; i++) {
         if (trainClass[i] == category) {
            //So document i is of interest to us.
            final ArrayList<Integer> topNWords = Summarizer.getTopNRows(V, i, n);
            for (int j = 0; j < topNWords.size(); j++) {
               topWordsAcrossDocs.add(topNWords.get(j));
            }
         }
      }
      final Integer[] allTopWords = new Integer[n];
      topWordsAcrossDocs.toArray(allTopWords);
      Arrays.sort(allTopWords);
      return Arrays.copyOf(allTopWords, n);
   }
   
   public int getNumCategories(final int[] actualClass, final int trainingSetSize) {
      final HashMap<Integer, Integer> map = new HashMap<Integer, Integer>();
      int numberOfCategories = 0;
      for (int i = 0; i < actualClass.length; i++) {
         if (map.get(actualClass[i]) == null) {
            numberOfCategories++;
            map.put(actualClass[i], i);
         }
      }
      
      return numberOfCategories;
   }
   
   @Override
   public double cost(final SimpleMatrix V, final SimpleMatrix W, final SimpleMatrix H) {
      // TODO Auto-generated method stub
      return NormOps.fastNormF(V.minus(W.mult(H)).getMatrix());
   }
   
   @Override
   public String costString() {
      // TODO Auto-generated method stub
      return "Frobenius Norm";
   }
   
   @Override
   public String toString() {
      return "TopWords NMF Tool";
   }
}
